source("3.2.1 EtaEst.R")
source("3.2.2 GetRatio.R")
source("3.2.3 GetMaxKGen.R")
source("3.2.4 GetK.R")
source("3.2.5 PMaxEstRough.R")
source("3.2.6 nLL.R")
source("3.2.7 MarginalRatio.R")

source("3.5.1 DR.R")
source("3.5.2 PropScoreEst.R")
source("3.5.3 DR_MyFunc.R")

DR.EE.Val = function(data, cnt, K, seed, max.step = 500, thres = 1e-4, 
                 mle,  outcome.type = "MLE",dr.est){
  # Output: theta, time, contrast
  # nll is just filled with 0---DR has nothing to do with nll
  # max.step = 5;thres = 1e-4; cnt <- count; outcome.type = "MLE" #DEBUG
  # dr.est <- dr.est.MLE
  
  # Model for E[L_k+1|A_k, L_k, X]
  eta.coef        = EtaEst(data, cnt, K)
  X.names         = paste("X",1:pL0,sep="")
  names(eta.coef) = c("gamma00","gamma01","gamma10","gamma11","gamma.",X.names)
  
  # Model for Propensity Score
  b.coef  = B.Est(data, cnt, K)
  
  # Model for the outcome
  p.y.func <- NA
  theta.mle = mle$theta.coef[c(1:2, 4:7)]
  theta.dr = dr.est$theta.coef[c(1:2, 4:7)]
  # phi.coef = runif(pL0+3,  -epsilon, epsilon) # check dimension
  phi.coef = mle$phi.coef
  gop.coef = mle$gop.coef[c(1, 4:7)]
  # source("3.5.3 DR_MyFunc.R")
  if (outcome.type=="MLE"){
    p.y.func <- get.y.prob.from.MLE(data, theta.coef = theta.mle,
                                    eta.coef,gop.coef,phi.coef,
                                    cnt, rep=500, seed)
  }else if(outcome.type=="logit"){
    p.y.func = y.logit.est(data, cnt, K)
  }else{
    stop("Outcome.type must be either MLE or logit")
  }
  
  
  
  obj.val <- dr.obj.value(data, theta.dr=theta.dr, theta.mle=theta.mle,
                                   eta.coef,gop.coef, 
                                   p.y.func, b.coef, 
                                   cnt, rep=500, seed,
                                   max.step = max.step, thres = thres)
  
  return(obj.val )
  
}


dr.obj.value <- function(data, theta.dr, theta.mle, eta.coef,gop.coef, 
                                p.y.func, b.coef, 
                                cnt, rep=500, seed,
                                max.step, thres){
  
  N    = nrow(data)
  Ak   = cbind(data[,(pL0+1):(pL0+K+1)])
  Lk   = cbind(data[,(pL0+K+2):(pL0+2*K+2)])
  Y = data[,"Y"]
  X.base = data[,1:pL0]
  dr.objective <- function(pars){
    # pars <- theta.coef # DEBUG
    U.m <- E.U.m <- array(NA, dim=c(N, K+1))
    
    LK = Lk[,(K+1)]; AK = Ak[,(K+1)]
    U.K.cov <- cbind(1-LK, LK, X.base)
    # U.m[, K+1] <- c(Y * exp(-t(t(U.K.cov) * AK) %*% pars))
    U.m[, K+1] <- c(Y * exp(-(U.K.cov* AK) %*% pars))
    
    # A t(t()) trick for multiplying one vector to each row
    # By default d1 * v1 is multiplying v1 to each column of v1
    #test code for the t(t()) trick
    # c1 <- c(1,2,3)
    # c2 <- c(4,5,6)
    # c3 <- c(7,8,9)
    # d1 <- data.frame(c1,c2,c3)
    # v1 <- c(1,2,3)
    # t(t(d1)*v1)
    #       c1 c2 c3
    # [1,]  1  8 21
    # [2,]  2 10 24
    # [3,]  3 12 27
    # E.U.m.func <- list()
    E.U.m.func <- function(A.mminus1.bar, Lm.bar){
      Ak.K.eq.0 <- cbind(A.mminus1.bar, rep(0, N))
      return(p.y.func(A= Ak.K.eq.0, L = Lm.bar))
    }
    
    A.mminus1.bar <- Ak[,1:K]
    Lm.bar <- Lk
    E.U.m[, K+1] <- E.U.m.func(A.mminus1.bar, Lm.bar)
    E.U.mplus1.1.1 <- E.U.mplus1.0.1 <- 
      E.U.mplus1.1.0 <- E.U.mplus1.0.0 <- array(NA,dim=c(N, K+1))
    E.U.mplus1.1.1[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(1,N)) , cbind(Lk[,1:K], rep(1,N)))
    E.U.mplus1.0.1[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(1,N)) , cbind(Lk[,1:K], rep(0,N)))
    E.U.mplus1.1.0[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(0,N)) , cbind(Lk[,1:K], rep(1,N)))
    E.U.mplus1.0.0[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(0,N)) , cbind(Lk[,1:K], rep(0,N)))  
    
    for (m in K:1){
      # m <- K #DEBUG
      # m <- 1 # DEBUG
      # print(m)
      Lm = Lk[,m]; Am = Ak[,m]
      U.m.cov <- cbind(1-Lm, Lm, X.base)
      # U.m[, m] <- U.m[, (m+1)] *  c(Y * exp(-t(t(U.m.cov) * Am) %*% pars))
      U.m[, m] <- U.m[, (m+1)] *  c(Y * exp(-(U.m.cov) * Am) %*% pars)
      
      get.m.bar <- function(df, m){
        if (m == 1){
          return(matrix(df[,1:1], ncol=1))
        }else if (m <1){
          return(NULL)
        }else{
          return(df[,1:m])
        }
      }
      A.mminus1.bar <- get.m.bar(Ak, m-1)
      A.mminus2.bar <- get.m.bar(Ak, m-2)
      L.mminus1.bar <- get.m.bar(Lk, m-1)
      Lm.bar <- get.m.bar(Lk, m)
      
      
      E.U.m.func <- function(A.mminus1.bar, Lm.bar){
        # A.mminus1.bar is of dimension c(N, m-1)
        # Lm.bar is of dimension c(N, m)
        Lm = Lm.bar[,ncol(Lm.bar)]
        p.Am <- c(prob.score.func(b.coef,L=Lm,X=X.base))
        U.m.cov <- cbind(1-Lm, Lm, X.base)
        exp.factor <- c(exp(-U.m.cov %*% theta.mle))
        
        p.Lm.1.1 <- p.L.func(eta.coef= eta.coef, ak= rep(1,N), lk= Lm, X=X.base)
        p.Lm.0.1 <- (1-p.Lm.1.1)
        p.Lm.1.0 <- p.L.func(eta.coef= eta.coef, ak= rep(0,N), lk= Lm, X=X.base)
        p.Lm.0.0 <- 1-p.Lm.1.0
        res <- (p.Am * exp.factor * 
                  (p.Lm.1.1*E.U.mplus1.1.1[, m+1] + p.Lm.0.1*E.U.mplus1.0.1[, m+1]) +
                  (1- p.Am) * 
                  (p.Lm.1.0*E.U.mplus1.1.0[, m+1] + p.Lm.0.0*E.U.mplus1.0.0[, m+1])
        )
        return(res)
      }
      
      E.U.m[, m] <- E.U.m.func(A.mminus1.bar=A.mminus1.bar,
                               Lm.bar = Lm.bar)
      
      E.U.mplus1.1.1[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(1,N)) , cbind(L.mminus1.bar, rep(1,N)))
      E.U.mplus1.0.1[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(1,N)) , cbind(L.mminus1.bar, rep(0,N)))
      E.U.mplus1.1.0[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(0,N)) , cbind(L.mminus1.bar, rep(1,N)))
      E.U.mplus1.0.0[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(0,N)) , cbind(L.mminus1.bar, rep(0,N)))  
    }
    
    d.m <- E.d.m <- array(NA, dim=c(N, pL0 + 2))
    total <- array(0, dim=c(N,pL0+2))
    for (m in 1:(K+1)){
      # m <- 1
      Lm <- Lk[,m]; Am <- Ak[,m]
      b.cov <- cbind(rep(1,N),Lm, X.base)
      p.Am <- c(prob.score.func(b.coef,L=Lm,X=X.base))
      
      d.m <- (b.cov)*Am
      E.d.m <- (b.cov)*p.Am
      total <- total + (d.m - E.d.m) * (U.m[,m] - E.U.m[,m])
      # d.m[,m,] <- t(t(b.cov)*Am)
      # E.d.m[,m,] <- t(t(b.cov)*p.Am)
    }
    
    tmp <- colSums(total * cnt)
    
    # for (m in 1:(K+1)){
    #   total <- total + (d.m[,m,] - E.d.m[,m,]) * (U.m[,m] - E.U.m[,m])
    # }
    
    return(sum(tmp^2))
  }
  pars <- theta.dr
  return(dr.objective(pars=pars))
}
